/**
 * @file main.cpp
 * @author SDK_TEAM
 * @brief
 * @version 0.1
 * @date 2022-10-26
 *
 * Copyright:
 * © 2018 北京灵汐科技有限公司 版权所有。
 * 注意：以下内容均为北京灵汐科技有限公司原创，未经本公司允许，不得转载，否则将视为侵权；对于不遵守此声明或者其他违法使用以下内容者，本公司依法保留追究权。
 * © 2018 Lynxi Technologies Co., Ltd. All rights reserved.
 * NOTICE: All information contained here is, and remains the property of Lynxi.
 * This file can not be copied or distributed without the permission of Lynxi
 * Technologies Co., Ltd.
 *
 */

#include <lyn_api.h>
#include <lyn_blas.h>
#include <nlohmann/json.hpp>
#include <string>
#include "imageDecoder.h"
#include "imageEncoder.h"
#include "ipeParamEx.h"
#include "osd_plugin.h"
#include "face_post_process.h"
#include "typeConv.hpp"
#include "util.hpp"
#include "CLI/CLI11.hpp"

int main(int argc, char *argv[])
{
  CLI::App app;

  std::string inputFile;
  std::string outputFile;
  int deviceId = 0;
  std::string faceDetectModel;
  std::string faceDetectPostPluginPath;
  std::string osdPluginPath;

  app.add_option("input_image", inputFile, "input jpeg file path")->default_val("../data/retinaface_face_detection_with_image_sample/test.jpeg");
  app.add_option("output_image", outputFile, "output jpeg file path")->default_val("../output.jpeg");
  app.add_option("device_id", deviceId, "lynxi device id")->default_val(0);
  app.add_option("-m,--face_detect_model", faceDetectModel, "face detect model path")->default_val("../model/retinaface_mb.25_b1_h640_w640_c3_iuint8_ofloat16_0000/Net_0");
  app.add_option("-p,--face_detect_post_plugin", faceDetectPostPluginPath, "face detect post process plugin path")->default_val("../common/plugin/postprocess/cross_compiled_lib/libRetinafacePostPlugin.so");
  app.add_option("-o,--osd_plugin", osdPluginPath, "osd plugin path")->default_val("../common/plugin/osd/lib/libOsdPlugin.so");

  CLI11_PARSE(app, argc, argv);

  if (!IsFileExist(inputFile))
  {
    std::cout << "error : input jpeg file " << inputFile << " is not exsit." << std::endl;
    return 0;
  }

  if (!IsValidDevice(deviceId))
  {
    std::cout << "error : device id is invalid." << std::endl;
    return 0;
  }

  if (!IsFileExist(faceDetectModel))
  {
    std::cout << "error : face detect model is not exsit." << std::endl;
    return 0;
  }

  if (!IsFileExist(faceDetectPostPluginPath))
  {
    std::cout << "error : face detect post process plugin is not exsit." << std::endl;
    return 0;
  }

  if (!IsFileExist(osdPluginPath))
  {
    std::cout << "error : osd plugin is not exsit." << std::endl;
    return 0;
  }

  // 2. 创建 context 与 stream
  lynContext_t ctx = nullptr;
  CHECK_ERR(lynCreateContext(&ctx, deviceId));
  CHECK_ERR(lynRegisterErrorHandler(StreamErrorHandler, nullptr));
  CHECK_ERR(lynCosineInit());
  lynStream_t stream = nullptr;
  CHECK_ERR(lynCreateStream(&stream));

  // 3. 初始化解码图片类
  std::cout << "image decode" << std::endl;
  ImageDecoder imageDec;
  imageDec.Init(inputFile, false);
  lynImageInfo_t imgInfo;
  imageDec.GetImageDecInfo(&imgInfo);

  // 4. 同步解码图片到 Device 侧
  lynCodecBuf_t decodeImg;
  decodeImg.size = imgInfo.output.predictBufSize;
  CHECK_ERR(lynMalloc((void **)&decodeImg.data, decodeImg.size));
  imageDec.DecodeImageToDevice(stream, decodeImg.data);

  // 5. 加载模型
  ModelInfo faceDetectModelInfo;
  faceDetectModelInfo.LoadModelByPath(faceDetectModel.c_str());

  // 加载lynxi plugin后处理和osd
  lynPlugin_t faceDetectPlugin;
  CHECK_ERR(
      lynPluginRegister(&faceDetectPlugin, faceDetectPostPluginPath.c_str()));
  lynPlugin_t osdPlugin;
  CHECK_ERR(lynPluginRegister(&osdPlugin, osdPluginPath.c_str()));

  // 6. 人脸检测IPE处理
  std::cout << "detect faces" << std::endl;
  uint8_t *pDetectIpeBuf = nullptr;
  CHECK_ERR(lynMalloc((void **)&pDetectIpeBuf, faceDetectModelInfo.inputSize));
  IpeParamFace ipeFace(faceDetectModelInfo.width, faceDetectModelInfo.height);
  ipeFace.SetImgInfo(imgInfo.output.width, imgInfo.output.height,
                     LYN_PIX_FMT_NV12);
  ipeFace.CalcParam(stream, decodeImg.data, pDetectIpeBuf);

  // 7. 调用 APU 推理接口
  uint8_t *pDevDetectApuBuf = nullptr;
  CHECK_ERR(
      lynMalloc((void **)&pDevDetectApuBuf, faceDetectModelInfo.outputSize));
  CHECK_ERR(lynExecuteModelAsync(stream, faceDetectModelInfo.model,
                                 pDetectIpeBuf, pDevDetectApuBuf,
                                 faceDetectModelInfo.batchSize));

  // 8. 加载人脸检测post Plugin, 进行模型后处理
  FacePostProcessInfo_t post_info;
  memset(&post_info, 0, sizeof(FacePostProcessInfo_t));

  post_info.score_threshold = 0.25;
  post_info.nms_threshold = 0.45;
  post_info.width = faceDetectModelInfo.width;
  post_info.height = faceDetectModelInfo.height;
  post_info.ori_width = imgInfo.output.width;
  post_info.ori_height = imgInfo.output.height;
  FaceDetectInfo *pDevDetectInfo;
  CHECK_ERR(lynMalloc((void **)&pDevDetectInfo, sizeof(FaceDetectInfo)));
  post_info.output_tensor = pDevDetectApuBuf;
  post_info.detectInfo = pDevDetectInfo;

  CHECK_ERR(lynPluginRunAsync(stream, faceDetectPlugin, "lynFacePostProcess",
                              &post_info, sizeof(post_info)));

  FaceDetectInfo *pHostDetectInfo = new FaceDetectInfo;
  CHECK_ERR(lynMemcpyAsync(stream, pHostDetectInfo, pDevDetectInfo,
                           sizeof(FaceDetectInfo), ServerToClient));
  lynSynchronizeStream(stream);

  lynBoxesInfo *pHostBoxesInfo = new lynBoxesInfo;
  memset(pHostBoxesInfo, 0, sizeof(lynBoxesInfo));
  pHostBoxesInfo->boxesNum = pHostDetectInfo->boxesNum;
  for (int i = 0; i < pHostBoxesInfo->boxesNum; ++i)
  {

    pHostBoxesInfo->boxes[i].xmin = pHostDetectInfo->boxes[i].xmin;
    pHostBoxesInfo->boxes[i].ymin = pHostDetectInfo->boxes[i].ymin;
    pHostBoxesInfo->boxes[i].xmax = pHostDetectInfo->boxes[i].xmax;
    pHostBoxesInfo->boxes[i].ymax = pHostDetectInfo->boxes[i].ymax;
    pHostBoxesInfo->boxes[i].id = pHostDetectInfo->boxes[i].id;
    pHostBoxesInfo->boxes[i].score = pHostDetectInfo->boxes[i].score;
  }

  // 9. 加载osd Plugin, 将结果叠加到图像上
  lynBoxesInfo *pDevBoxesInfo;
  CHECK_ERR(lynMalloc((void **)&pDevBoxesInfo, sizeof(lynBoxesInfo)));
  CHECK_ERR(lynMemcpyAsync(stream, pDevBoxesInfo, pHostBoxesInfo,
                           sizeof(lynBoxesInfo), ClientToServer));

  lynDrawBoxAndTextPara para;
  para.imgData = decodeImg.data;
  para.imgFmt = LYN_PIX_FMT_NV12;
  para.imgW = imgInfo.output.width;
  para.imgH = imgInfo.output.height;
  para.boxesInfo = pDevBoxesInfo;
  para.boxColor = DRAW_COLOR_BLUE;
  para.boxThick = DRAW_THICK_2;
  para.fontSize = FONT_SIZE_24;
  para.fontColor = DRAW_COLOR_BLUE;
  CHECK_ERR(lynPluginRunAsync(stream, osdPlugin, "lynDrawBoxAndText", &para,
                              sizeof(para)));
  SetCallback(stream, [pDevBoxesInfo]
              { CHECK_ERR(lynFree(pDevBoxesInfo)); });

  // 9. 进行图片编码
  std::cout << "save new image" << std::endl;
  ImageEncoder encoder;
  encoder.EncodeImageAndSave(stream, outputFile, decodeImg, imgInfo);

  // 10. 销毁资源
  CHECK_ERR(lynDestroyStream(stream));
  CHECK_ERR(lynPluginUnregister(faceDetectPlugin));
  CHECK_ERR(lynPluginUnregister(osdPlugin));
  CHECK_ERR(lynFree(decodeImg.data));
  CHECK_ERR(lynFree(pDetectIpeBuf));
  CHECK_ERR(lynFree(pDevDetectApuBuf));
  // if (pDevFeatureApuBuf) CHECK_ERR(lynFree(pDevFeatureApuBuf));
  if (pHostBoxesInfo)
    delete pHostBoxesInfo;
  faceDetectModelInfo.UnLoadModel();
  CHECK_ERR(lynCosineUninit());
  CHECK_ERR(lynDestroyContext(ctx));

  std::cout << "process over" << std::endl;

  return 0;
}
